"""
PCD 转连通节点图一体化处理脚本
@author: CeasarSmj
@date: 2025-5-16
"""

import os
import torch
import numpy as np
from scipy.spatial import KDTree
import yaml
import open3d as o3d
import matplotlib.pyplot as plt

# ================== 参数配置 ==================
CONFIG = {
    # 输入输出路径（硬编码）
    "input_pcd": "input.pcd",          # 输入PCD文件路径
    "output_graph": "output.yaml",     # 输出图文件路径
    
    # 体素地图参数
    "voxel": {
        "raw_resolution": 0.05,          # 原始体素分辨率
        "target_resolution": 0.20,       # 目标体素分辨率
        "gaussian_kernel": 100,          # 高斯滤波核大小
        "gaussian_sigma": 1,             # 高斯分布标准差
        "runnable_radius": 2,            # 可行域检测半径
        "runnable_threshold": 0,         # 可行域阈值
        "enable_visualization": False    # 是否启用中间可视化
    },
    # 节点图参数
    "graph": {
        "block_size": 4.0,               # 体素块大小
        "connect_dist": 0.3              # 块间连接距离阈值
    }
}
# ===========================================

class VoxelMapGenerator:
    """体素地图生成器（来自pcd_to_runnable.py）"""
    def __init__(self, config):
        self.config = config["voxel"]
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        for k, v in self.config.items():
            setattr(self, k, v)
        self._init_kernel()
    
    def _init_kernel(self):
        """初始化卷积核"""
        self.kernel_radius_N = max(1, int(self.runnable_radius / self.target_resolution))
        def weight_func(r, h):
            return np.exp(-(r**2)) * ((1/(1+np.exp(h)))*4 - (1/(1+np.exp(-h)))*2 - np.exp(-(h**2))*5)
        self.runnable_kernel = torch.tensor([[[
            weight_func(np.sqrt(x**2 + y**2), z)
            for z in range(-self.kernel_radius_N, self.kernel_radius_N + 1)
        ] for y in range(-self.kernel_radius_N, self.kernel_radius_N + 1)]
        for x in range(-self.kernel_radius_N, self.kernel_radius_N + 1)]], device=self.device)

    def generate_from_pcd(self, file_path):
        """从PCD生成可行域体素地图"""
        points = self.load_pcd(file_path)
        raw_voxel, origin = self.generate_voxel(points)
        filtered = self.filter_voxel(raw_voxel)
        runnable = self.create_runnable_map(filtered)
        cropped, transform = self.crop_and_transform(origin)
        return {
            "voxel": cropped,
            "transform": transform,
            "resolution": self.target_resolution
        }

    def load_pcd(self, file_path):
        """加载PCD文件"""
        pcd = o3d.io.read_point_cloud(file_path)
        return np.asarray(pcd.points)

    def generate_voxel(self, points):
        """生成原始体素地图"""
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(
            pcd, self.raw_resolution)
        voxels = voxel_grid.get_voxels()
        if not voxels:
            return np.zeros((0,0,0)), np.zeros(3)
        indices = np.array([v.grid_index for v in voxels])
        min_indices = indices.min(axis=0)
        max_indices = indices.max(axis=0)
        shape = max_indices - min_indices + 1
        voxel_array = np.zeros(shape, dtype=np.uint8)
        adjusted_indices = indices - min_indices
        voxel_array[tuple(adjusted_indices.T)] = 1
        pcd_min = points.min(axis=0)
        origin = pcd_min + min_indices * self.raw_resolution
        return voxel_array, origin

    def filter_voxel(self, voxel_map):
        """滤波处理"""
        kernel = self.create_gaussian_kernel_3d()
        tensor = torch.from_numpy(voxel_map).float().to(self.device).unsqueeze(0).unsqueeze(0)
        padding = self.gaussion_kernel_size // 2
        filtered = torch.nn.functional.conv3d(tensor, kernel.unsqueeze(0).unsqueeze(0), padding=padding)
        threshold = torch.max(filtered) * 0.05
        filtered = (filtered > threshold).float()
        scale_factor = self.target_resolution / self.raw_resolution
        kernel_size = int(scale_factor)
        filtered = torch.nn.functional.avg_pool3d(filtered, kernel_size, kernel_size)
        filtered = (filtered > 0.01).float()
        return filtered.squeeze().cpu().numpy().astype(np.uint8)

    def create_gaussian_kernel_3d(self):
        """创建3D高斯核"""
        kernel_size = self.gaussion_kernel_size
        sigma = self.gaussion_sigma
        x = torch.linspace(-2, 2, kernel_size, device=self.device)
        y = torch.linspace(-2, 2, kernel_size, device=self.device)
        z = torch.linspace(-2, 2, kernel_size//2, device=self.device)
        x, y, z = torch.meshgrid(x, y, z, indexing="ij")
        kernel = torch.exp(-(x.pow(2)+y.pow(2)+z.pow(2))/(2*sigma**2))
        return kernel / kernel.sum()

    def create_runnable_map(self, voxel_map):
        """创建可行域地图"""
        tensor = torch.from_numpy(voxel_map).float().to(self.device)
        conv1 = torch.nn.functional.conv3d(
            tensor.unsqueeze(0).unsqueeze(0), 
            self.runnable_kernel.unsqueeze(0).unsqueeze(0), 
            padding=self.kernel_radius_N
        )
        runnable = (conv1 > self.runnable_threshold).float().squeeze()
        conv2 = torch.zeros_like(runnable)
        if runnable.size(2) > 1:
            conv2[:, :, :-1] = runnable[:, :, 1:]
        mask = (runnable > 0) & (conv2 > 0)
        result = torch.zeros_like(runnable)
        result[mask] = 1
        return result.cpu().numpy().astype(np.uint8)

    def crop_and_transform(self, origin):
        """裁剪并计算坐标变换"""
        non_zero = np.argwhere(self.voxel_map)
        if non_zero.size == 0:
            return self.voxel_map, (0,0,0,0,0,0)
        min_coords = non_zero.min(axis=0)
        max_coords = non_zero.max(axis=0)
        cropped = self.voxel_map[
            min_coords[0]:max_coords[0]+1,
            min_coords[1]:max_coords[1]+1,
            min_coords[2]:max_coords[2]+1
        ]
        translation = origin + min_coords * self.target_resolution
        return cropped, tuple(translation) + (0, 0, 0)

class GraphBuilder:
    """节点图构建器（来自runnable_to_graph.py）"""
    def __init__(self, config):
        self.config = config["graph"]
        for k, v in self.config.items():
            setattr(self, k, v)

    def build_from_voxel(self, voxel_data, resolution, transform):
        """从体素数据构建节点图"""
        actual_points = self._convert_to_world_coords(voxel_data, resolution, transform)
        block_list = self._cluster_blocks(actual_points)
        nodes = self._create_nodes(block_list)
        edges = self._create_edges(block_list)
        return {"nodes": nodes, "edges": edges}

    def _convert_to_world_coords(self, voxel_data, resolution, transform):
        """将体素坐标转换为世界坐标"""
        points = np.argwhere(voxel_data)
        return points * resolution + np.array(transform[:3])

    def _cluster_blocks(self, points):
        """体素聚类"""
        n_voxels = len(points)
        if n_voxels == 0:
            return []
        mask = np.zeros(n_voxels, dtype=bool)
        block_list = []
        while not np.all(mask):
            remaining_indices = np.where(~mask)[0]
            if not remaining_indices.size:
                break
            seed_idx = np.random.choice(remaining_indices)
            current_block = [points[seed_idx]]
            mask[seed_idx] = True
            for _ in range(3):  # 最多扩展3次
                added = False
                center = np.mean(current_block, axis=0)
                remaining_indices = np.where(~mask)[0]
                if not remaining_indices.size:
                    break
                remaining_voxels = points[remaining_indices]
                dist_to_center = np.linalg.norm(remaining_voxels - center, axis=1)
                in_radius_mask = dist_to_center <= (self.block_size / 2)
                in_radius_voxels = remaining_voxels[in_radius_mask]
                in_radius_indices = remaining_indices[in_radius_mask]
                if not in_radius_voxels.size:
                    break
                tree = KDTree(current_block)
                dists, _ = tree.query(in_radius_voxels, k=1)
                connected_mask = dists <= self.connect_dist
                connected_voxels = in_radius_voxels[connected_mask]
                connected_indices = in_radius_indices[connected_mask]
                if connected_voxels.size > 0:
                    current_block.extend(connected_voxels.tolist())
                    mask[connected_indices] = True
                    added = True
            block_list.append(np.array(current_block))
        return block_list

    def _create_nodes(self, block_list):
        """创建节点"""
        nodes = []
        for i, block in enumerate(block_list):
            center = np.mean(block, axis=0)
            nodes.append({
                "id": i,
                "x": float(center[0]),
                "y": float(center[1]),
                "z": float(center[2])
            })
        return nodes

    def _create_edges(self, block_list):
        """创建边"""
        edges = []
        edge_id = 0
        block_trees = [KDTree(block) for block in block_list]
        is_isolated = [True for _ in range(len(block_list))]
        
        for i in range(len(block_list)):
            for j in range(i + 1, len(block_list)):
                distances, _ = block_trees[i].query(block_list[j], k=1)
                if np.any(distances <= self.connect_dist):
                    edges.append({
                        "id": edge_id, 
                        "from": i, 
                        "to": j
                    })
                    edge_id += 1
                    is_isolated[i] = False
                    is_isolated[j] = False
                    
        return [self.graph["nodes"][i] for i in range(len(self.graph["nodes"])) if not is_isolated[i]]

def main():
    """主处理流程"""
    # 1. 生成体素地图
    voxel_gen = VoxelMapGenerator(CONFIG)
    voxel_result = voxel_gen.generate_from_pcd(CONFIG["input_pcd"])
    
    # 2. 构建节点图
    graph_builder = GraphBuilder(CONFIG)
    graph = graph_builder.build_from_voxel(
        voxel_result["voxel"],
        voxel_result["resolution"],
        voxel_result["transform"]
    )
    
    # 3. 保存结果
    with open(CONFIG["output_graph"], 'w') as f:
        yaml.dump(graph, f)
    print(f"节点图已保存至 {CONFIG['output_graph']}")

if __name__ == "__main__":
    main()